package mr

import (
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/rpc"
	"os"
	"sort"
)

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

func write_intermediate(kva []KeyValue, reply *AssnJobReply) {
	encoders := []*json.Encoder{}
	file_names := []string{}
	func ()  {
		for i:=0; i<reply.Nreduce; i++ {
			file_handle, err := os.CreateTemp(".", "map_for_reduce_tmp_*")
			if err != nil {
				log.Fatalf("cannot create tmp file map_for_reduce_tmp_*")
			}
			encoders = append(encoders, json.NewEncoder(file_handle))
			file_names = append(file_names, file_handle.Name())
			defer file_handle.Close()
		}
		for _, kv := range(kva) {
			err := encoders[ihash(kv.Key)%reply.Nreduce].Encode(&kv)
			if err != nil {
				log.Fatalf("cannot encode %v", kv)
			}
		}
	}()

	for i:=0; i<len(file_names); i++ {
		err := os.Rename(file_names[i], fmt.Sprintf("mr-%d-%d", reply.JobIndex, i))
		if err != nil {
			log.Fatalf("cannot rename file %v", file_names[i])
		}
	}
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {
	for {
		reply := AssnJobReply{}
		if !callAssnJob(&reply) {
			return
		}
		if reply.JobIndex == -1 {
			continue
		}

		if reply.IsMapJob {
			file, err := os.Open(reply.Inputfiles[0])
			if err != nil {
				log.Fatalf("cannot open %v", reply.Inputfiles[0])
			}
			content, err := io.ReadAll(file)
			if err != nil {
				log.Fatalf("cannot read %v", reply.Inputfiles[0])
			}
			file.Close()
			kva := mapf(reply.Inputfiles[0], string(content))
			write_intermediate(kva, &reply)
			if !callFinishJob(reply.JobIndex, true) {
				return
			}
		} else {
			kva := []KeyValue{}
			for _, file:=range(reply.Inputfiles) {
				file_handle, err := os.Open(file)
				if err != nil {
					log.Fatalf("cannot open %v", file)
				}
				dec := json.NewDecoder(file_handle)
				for {
					var kv KeyValue
					if err := dec.Decode(&kv); err != nil {
						break
					}
					kva = append(kva, kv)
				}
			}
			sort.Sort(ByKey(kva))

			ofile, err := os.CreateTemp(".", "reduce_out_*")
			tmpfile_name := ofile.Name()
			if err != nil {
				log.Fatalf("cannot create tmp file %v", tmpfile_name)
			}
			i := 0
			for i < len(kva) {
				j := i + 1
				for j < len(kva) && kva[j].Key == kva[i].Key {
					j++
				}
				values := []string{}
				for k := i; k < j; k++ {
					values = append(values, kva[k].Value)
				}
				output := reducef(kva[i].Key, values)

				// this is the correct format for each line of Reduce output.
				fmt.Fprintf(ofile, "%v %v\n", kva[i].Key, output)

				i = j
			}
			ofile.Close()
			os.Rename(tmpfile_name, fmt.Sprintf("mr-out-%d", reply.JobIndex))
			if !callFinishJob(reply.JobIndex, false) {
				return
			}
		}
	}
}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func callAssnJob(reply *AssnJobReply) bool {
	args := AssnJobArgs{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := false
	cnt := 0
	for !ok {
		ok = call("Coordinator.AssnJob", &args, reply)
		cnt++
		if cnt >= 3 {
			break
		}
	}
	if ok {
		// fmt.Printf("AssnJobReply.IsMapJob %v\n", reply.IsMapJob)
		return true
	}
	fmt.Printf("call failed!\n")
	return false
}

func callFinishJob(jobIndex int, isMapJob bool) bool {
	args := FinishJobArgs{}
	reply := FinishJobReply{}
	args.JobIndex=jobIndex
	args.IsMapJob=isMapJob
	ok := false
	cnt := 0
	for !ok {
		ok = call("Coordinator.FinishJob", &args, &reply)
		cnt++
		if cnt >= 3 {
			break
		}
	}
	if ok {
		return true
	}

	fmt.Printf("call failed!\n")
	return false
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
